#!-!encoding=utf-8!-!
import sys

import torch
from tqdm import tqdm
from transformers import AlbertModel, BertTokenizer

from medical_ner.models.model import build_model
from medical_ner.predictor import Predictor
from medical_ner.utils.constants import ID2TAG
from medical_ner.utils.dataloader import NERDataLoader
from medical_ner.utils.logging_utils import set_logger
from medical_ner.utils.metrics import accuracy, ner_accuracy, get_entities
from medical_ner.utils.param import Param, build_params

print(sys.path)


def t0():
    param = Param()
    model = build_model(param)
    print(model)

    input_ids = torch.tensor([
        [1, 2, 3, 4, 5],
        [2, 5, 6, 9, 0]
    ], dtype=torch.long)
    attention_mask = torch.tensor([
        [1.0, 1.0, 1.0, 1.0, 1.0],
        [1.0, 1.0, 1.0, 1.0, 0.0],
    ])

    r = model(input_ids, attention_mask)
    print(r)

    tokenizer_class = model.lm_model.tokenizer_class
    print(tokenizer_class)
    tokenizer = tokenizer_class.from_pretrained(
        r'C:\Users\HP\dataroot\models\bert-base-chinese'
    )
    print(tokenizer)
    print(tokenizer("我爱游泳"))

    print("=" * 100)
    for name, param in model.named_parameters():
        if param.requires_grad:
            print(name, "=" * 3, param.shape)

    # torch.onnx.export(
    #     model.eval(),
    #     (input_ids, attention_mask),
    #     'model.onnx',
    #     input_names=['token_ids', 'token_mask'],
    #     output_names=['scores'],  # 由于此时返回的结果中loss和output为None，所以实际返回只有一个
    #     opset_version=12,
    #     dynamic_axes={
    #         'token_ids': {
    #             0: 'n',
    #             1: 't'
    #         },
    #         'token_mask': {
    #             0: 'n',
    #             1: 't'
    #         },
    #         'scores': {
    #             0: 'n',
    #         },
    #     }
    # )


def t1():
    from transformers import BertModel
    from transformers import BertConfig

    # python model_download.py --repo_id bert-base-chinese
    # NOTE: 当from_pretrained参数为模型id的时候，模型会从huggingface.co这个网站下载对应的模型文件，并放到C:\Users\HP\.cache\huggingface\hub
    # 手动下载，通过https://aliendao.cn/进行下载: python model_download.py --repo_id bert-base-chinese
    # 迁移的时候直接给定路径即可
    # bert = BertModel.from_pretrained("bert-base-chinese")
    bert = BertModel.from_pretrained(
        pretrained_model_name_or_path=r"C:\Users\HP\dataroot\models\bert-base-chinese"
    )
    # bert = BertModel(config=BertConfig(), add_pooling_layer=False)  # 完全重新创建
    print(bert)
    print(next(bert.parameters()).mean())
    print(next(bert.parameters())[:3, :3])

    input_ids = torch.tensor([
        [1, 2, 3, 4, 5],
        [2, 5, 6, 9, 0]
    ], dtype=torch.long)
    attention_mask = torch.tensor([
        [1.0, 1.0, 1.0, 1.0, 1.0],
        [1.0, 1.0, 1.0, 1.0, 0.0],
    ])

    outputs = bert.forward(
        input_ids=input_ids,
        attention_mask=attention_mask,
        output_hidden_states=False,  # 是否返回所有层的hidden_states
        return_dict=False  # 以对象的形式返回，默认就是True
    )
    print(outputs)

    torch.onnx.export(
        bert,
        (input_ids, attention_mask),
        'bert.onnx',
        input_names=['input_ids', 'attention_mask'],
        output_names=['token_feat'],
        opset_version=12,
        dynamic_axes={
            'input_ids': {
                0: 'n',
                1: 't'
            },
            'attention_mask': {
                0: 'n',
                1: 't'
            },
            'token_feat': {
                0: 'n',
                1: 't'
            },
        }
    )


def t2():
    from transformers import AlbertModel, AlbertConfig, AlbertTokenizer
    config = AlbertConfig(
        num_hidden_layers=4,  # 共享的group 层有多少层
        num_hidden_groups=1,  # 整个layer有多少个group
        inner_group_num=3  # 每个group中有几层(group内的参数是不共享的)
    )
    # bert = AlbertModel.from_pretrained("albert-base-v1")
    bert = AlbertModel(config=config, add_pooling_layer=False)  # 完全重新创建
    print(bert)

    input_ids = torch.tensor([
        [1, 2, 3, 4, 5],
        [2, 5, 6, 9, 0]
    ], dtype=torch.long)
    attention_mask = torch.tensor([
        [1.0, 1.0, 1.0, 1.0, 1.0],
        [1.0, 1.0, 1.0, 1.0, 0.0],
    ])

    # noinspection PyTypeChecker
    outputs = bert.forward(
        input_ids=input_ids,
        attention_mask=attention_mask,
        output_hidden_states=False,  # 是否返回所有层的hidden_states
        return_dict=False  # 以对象的形式返回，默认就是True
    )
    print(outputs)

    # torch.onnx.export(
    #     bert,
    #     (input_ids, attention_mask),
    #     'albert.onnx',
    #     input_names=['input_ids', 'attention_mask'],
    #     output_names=['token_feat'],
    #     opset_version=12,
    #     dynamic_axes={
    #         'input_ids': {
    #             0: 'n',
    #             1: 't'
    #         },
    #         'attention_mask': {
    #             0: 'n',
    #             1: 't'
    #         },
    #         'token_feat': {
    #             0: 'n',
    #             1: 't'
    #         },
    #     }
    # )


def t3():
    from medical_ner.utils.dataloader import read_examples
    from medical_ner.utils.dataloader import convert_example_to_feature

    tokenizer = BertTokenizer.from_pretrained(
        pretrained_model_name_or_path=r"C:\Users\HP\dataroot\models\bert-base-chinese"
    )
    examples = read_examples(
        r"..\datas\sentence_tag_small\test.txt"
    )
    features = convert_example_to_feature(examples, tokenizer, 256)
    print(len(features))
    print(features[0].input_ids)
    print(features[0].label_ids)


def t4():
    param = Param()
    model = build_model(param)
    print(model)

    ner = NERDataLoader(
        tokenizer_class=BertTokenizer,
        param=param
    )
    dataloader = ner.get_dataloader("train")
    for input_ids, attention_mask, label_ids in dataloader:
        print(input_ids.shape)
        scores, loss, pred_label_ids = model(input_ids, attention_mask, labels=label_ids, return_output=True)
        print(scores.shape)
        print(label_ids.shape)
        print(loss)
        print(pred_label_ids.shape)
        acc1 = accuracy(pred_label_ids, label_ids, attention_mask)
        acc2 = ner_accuracy(pred_label_ids, label_ids)
        print(acc1)
        print(acc2)
        print(label_ids)
        break


def t5():
    from medical_ner import training

    training.run()


def t6():
    param = Param()
    model = build_model(param).to(param.device)
    ckpt = torch.load(r"outputs\models\ckpt\000004.pkl", map_location="cpu")
    # 参数恢复
    missing_keys, unexpected_keys = model.load_state_dict(ckpt['param'], strict=False)
    print(missing_keys, unexpected_keys)

    dataloader = NERDataLoader(model.lm_model.tokenizer_class, param)
    test_loader = dataloader.get_dataloader("test")
    bar = tqdm(test_loader)

    model.eval()
    for input_ids, attention_mask, label_ids in bar:
        scores, loss, batch_pred_label_ids = model(input_ids, attention_mask, labels=label_ids, return_output=True)
        print(batch_pred_label_ids)
        print(label_ids)
        # 恢复真实标签的信息
        batch_size, t = attention_mask.shape
        for i in range(batch_size):
            real_len = int(attention_mask[i].sum().item())
            real_label_ids = label_ids[i][:real_len].to('cpu').numpy()
            pred_label_ids = batch_pred_label_ids[i][:real_len].to('cpu').numpy()
            true_tags = []
            pred_tags = []
            for real_label_id, pred_label_id in zip(real_label_ids, pred_label_ids):
                true_tags.append(ID2TAG[real_label_id])
                pred_tags.append(ID2TAG[pred_label_id])
            print("=" * 100)
            print(true_tags)
            print(pred_tags)
            true_entities = set(get_entities(true_tags, False))
            pred_entities = set(get_entities(pred_tags, False))
            print(true_entities)
            print(pred_entities)
        break


def t7():
    set_logger()
    param = build_params(resume=True)
    pred = Predictor(param)

    text = "， 患 者 2 0 0 8 年 9 月 3 日 因 “ 腹 胀 ， 发 现 腹 部 包 块 ” 在 我 院 腹 科 行 手 术 探 查 ， 术 中 见 盆 腹 腔 肿 物 ， 与 肠 管 及 子 宫 关 系 密 切 ， 遂 行 “ 全 子 宫 左 附 件 切 除 + 盆 腔 肿 物 切 除 + 右 半 结 肠 切 除 + D I X O N 术 ” ， 术 后 病 理 示 颗 粒 细 胞 瘤 ， 诊 断 为 颗 粒 细 胞 瘤 I I I C 期 ， 术 后 自 2 0 0 8 年 1 1 月 起 行 B E P 方 案 化 疗 共 4 程 ， 末 次 化 疗 时 间 为 2 0 0 9 年 3 月 2 6 日 。 之 后 患 者 定 期 复 查 ， 2 0 1 5 - 6 - 1 ， 复 查 C T 示 ： 髂 嵴 水 平 上 腹 部 L 5 腰 椎 前 见 软 组 织 肿 块 ， 大 小 约 3 0 M M × 4 5 M M ， 密 度 欠 均 匀 ， 边 界 尚 清 楚 ， 轻 度 强 化 。 查 肿 瘤 标 志 物 均 正 常 。 于 2 0 1 5 - 7 - 6 行 剖 腹 探 查 + 膀 胱 旁 肿 物 切 除 + 骶 前 肿 物 切 除 + 肠 表 面 肿 物 切 除 术 ， 术 程 顺 利 ， ， 术 后 病 理 示 ： 膀 胱 旁 肿 物 及 骶 前 肿 物 符 合 颗 粒 细 胞 瘤 。 于 2 0 1 5 - 7 - 1 3 、 8 - 1 4 给 予 泰 素 2 4 0 M G + 伯 尔 定 6 0 0 M G 化 疗 2 程 ， 过 程 顺 利 。 出 院 至 今 ， 无 发 热 ， 无 腹 痛 、 腹 胀 ， 有 脱 发 ， 现 返 院 复 诊 ， 拟 行 再 次 化 疗 收 入 院 。 起 病 以 来 ， 精 神 、 胃 纳 、 睡 眠 可 ， 大 小 便 正 常 ， 体 重 无 明 显 改 变 。"
    # text = "， 患 者 2 0 0 8 年 9 月 3 日 因 “ 腹 胀 ， 发 现 腹 部 包 块 ” 在 我 院 腹 科 行 手 术 探 查 "
    text = "， 患 者 因 罹 患 “ 胃 癌 ” 于 2 0 1 3 - 1 0 - 2 9 在 我 院 予 行 全 麻 上 胃 癌 根 治 术 ， ， 术 中 见 ： 腹 腔 内 腹 水 ， 腹 膜 无 转 移 ， 肝 脏 未 触 及 明 显 转 移 性 灶 ， 肿 瘤 位 于 胃 体 、 胃 底 部 ， 小 弯 侧 偏 后 壁 ， 约 5 * 4 * 2 C M 大 小 ， 肿 瘤 已 侵 达 浆 膜 外 ， 第 1 、 3 组 淋 巴 结 肿 大 ， 肿 瘤 尚 能 活 动 ， 经 探 查 决 定 行 全 胃 切 除 ， 空 肠 J 字 代 胃 术 。 手 术 顺 利 ， 术 后 积 极 予 相 关 对 症 支 持 治 疗 ； ， 后 病 理 示 ： 胃 底 、 体 小 弯 侧 低 分 化 腺 癌 ， 部 分 为 印 戒 细 胞 癌 图 像 ， 蕈 伞 型 ， 面 积 5 . 2 * 3 . 5 C M ， 局 部 侵 达 粘 膜 上 层 ， 并 于 少 数 腺 管 内 查 见 癌 栓 。 双 端 切 线 及 另 送 “ 近 端 切 线 ” 未 查 见 癌 。 呈 三 组 （ 5 / 1 3 个 ） 淋 巴 结 癌 转 移 。 一 组 （ 7 个 ） 、 四 组 （ 1 3 个 ） 、 五 组 （ 1 个 ） 、 六 组 （ 4 个 ） 淋 巴 结 未 查 见 癌 。 ， 癌 组 织 免 疫 组 化 染 色 示 ： E R C C 1 （ + ） 、 β - T U B U L I N - I I I （ + ） 、 T S （ - ） 、 R R M 1 （ - ） 、 T O P O I I 阳 性 细 胞 数 约 2 0 % 、 C E R B B - 2 （ 2 + ） 🚗 。 依 据 患 者 病 情 及 肿 瘤 病 理 与 分 期 继 续 术 后 辅 助 性 化 疗 指 征 存 在 ， 患 者 及 家 属 拒 绝 化 疗 。 自 术 后 出 院 以 来 ， 患 者 一 般 情 况 保 持 良 好 ； 无 发 热 ， 偶 有 恶 心 ， 无 呕 吐 ， 无 反 酸 、 嗳 气 ， 无 明 显 进 食 不 适 ， 偶 有 进 食 后 轻 微 腹 胀 ， 无 腹 痛 。 现 患 者 为 行 进 一 步 复 查 并 必 要 时 适 当 处 理 而 再 来 我 院 就 诊 ， 门 诊 依 情 以 “ 胃 恶 性 肿 瘤 术 后 ” 收 入 院 。 目 前 患 者 精 神 及 情 绪 状 态 良 好 ， 食 欲 较 术 前 明 显 减 少 ， 饮 食 可 ， 夜 间 睡 眠 后 ； 今 8 个 月 体 重 减 轻 1 8 K G 。"
    text = "".join(text.split(" "))
    r = pred.predict(text)
    print(r)


if __name__ == '__main__':
    t7()
